Count vowels permutation [Matrix Exponentiation]

Time: O(LogN); Space: O(1); hard

Given an integer n, your task is to count how many strings of length n can be formed under the following rules:

  • Each character is a lower case vowel (‘a’, ‘e’, ‘i’, ‘o’, ‘u’)

  • Each vowel ‘a’ may only be followed by an ‘e’.

  • Each vowel ‘e’ may only be followed by an ‘a’ or an ‘i’.

  • Each vowel ‘i’ may not be followed by another ‘i’.

  • Each vowel ‘o’ may only be followed by an ‘i’ or a ‘u’.

  • Each vowel ‘u’ may only be followed by an ‘a’.

Since the answer may be too large, return it modulo 10^9 + 7.

Example 1:

Input: n = 1

Output: 5

Explanation:

  • All possible strings are: “a”, “e”, “i” , “o” and “u”.

Example 2:

Input: n = 2

Output: 10

Explanation:

  • All possible strings are: “ae”, “ea”, “ei”, “ia”, “ie”, “io”, “iu”, “oi”, “ou” and “ua”.

Example 3:

Input: n = 5

Output: 68

Constraints:

  • 1 <= n <= 2 * 10^4

Hints:

  1. Use dynamic programming.

  2. Let dp[i][j] be the number of strings of length i that ends with the j-th vowel.

  3. Deduce the recurrence from the given relations between vowels.

1. Dynamic programming (Matrix Exponentiation) [O(LogN), O(1)]

[3]:

class Solution1(object):
    """
    Time: O(LogN)
    Space: O(1)
    """
    def countVowelPermutation(self, n):
        """
        :type n: int
        :rtype: int
        """
        def matrix_mult(m1, m2):
            return [[sum(a * b for a, b in zip(m1_row, m2_col)) %MOD for m2_col in zip(*m2)] for m1_row in m1]

        def matrix_expo(m, pow):
            assert pow >= 0 and int(pow) == pow, "Only non-negative, integer powers allowed"
            result = [[int(i==j) for j in range(len(m))] \
                                 for i in range(len(m))]
            while pow:
                if pow % 2:
                    result = matrix_mult(result, m)
                m = matrix_mult(m, m)
                pow //= 2
            return result

        MOD = 10**9 + 7
        T = [[0, 1, 1, 0, 1],
             [1, 0, 1, 0, 0],
             [0, 1, 0, 1, 0],
             [0, 0, 1, 0, 0],
             [0, 0, 1, 1, 0]]

        return sum(map(sum, matrix_expo(T, n-1))) % MOD

[4]:

s = Solution1()

n = 1
assert s.countVowelPermutation(n) == 5

n = 2
assert s.countVowelPermutation(n) == 10

n = 5
assert s.countVowelPermutation(n) == 68

2. [O(N), O(1)]

[5]:

class Solution2(object):
    """
    Time:  O(n)
    Space: O(1)
    """
    def countVowelPermutation(self, n):
        """
        :type n: int
        :rtype: int
        """
        MOD = 10**9 + 7
        a, e, i, o, u = 1, 1, 1, 1, 1

        for _ in range(1, n):
            a, e, i, o, u = (e+i+u) % MOD, (a+i) % MOD, (e+o) % MOD, i, (i+o) % MOD

        return (a+e+i+o+u) % MOD

[6]:

s = Solution2()

n = 1
assert s.countVowelPermutation(n) == 5

n = 2
assert s.countVowelPermutation(n) == 10

n = 5
assert s.countVowelPermutation(n) == 68